import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

import gradio as gr

from utils import zip_folder, SupportedModels
from translator import PDFTranslator

def translate_pdf(model_type, model, target_language, filepath):
    output_file_path = os.path.splitext(filepath)[0]
    translator = PDFTranslator(model_type, model)
    translator.translate_pdf(filepath, 'markdown', output_file_path, target_language)
    zip_file_path = zip_folder(output_file_path, os.path.dirname(filepath))
    if zip_file_path:
        return zip_file_path
    else:
        gr.Error('Failed to create zip file')


if __name__ == "__main__":
    supported_models = SupportedModels()
    supported_models.initialize('models.yaml')

    all_supported_models = supported_models.get_all_models()
    all_supported_model_types = supported_models.get_all_model_types()

    supported_languages = ['Chinese', 'Japanese', 'Korean']

    translator_gui = gr.Interface(
        fn=translate_pdf,
        title='Mini AI Translator v2.0',
        inputs=[
            gr.Dropdown(
                all_supported_model_types, label='model_type',
                info='Type of the translation model to use (e.g., OpenAI, ZhipuAI, Tongyi).'
            ),
            gr.Dropdown(
                all_supported_models, label='model',
                info='The specific model to use for translation. For example, "gpt-3.5-turbo" for OpenAI or "glm-4" for ZhipuAI or "qwen-turbo" for Tongyi.'
            ),
            gr.Dropdown(
                supported_languages, label='target_language', info='Choose the target language for translation.'
            ),
            gr.File(label='Upload the PDF file that needs to be translated'),
        ],
        outputs=[
            gr.File(label='Download the translated file zip')
        ],
        examples=[
            ['OpenAI', 'gpt-3.5-turbo', 'Chinese', None],
            ['ZhipuAI', 'glm-4', 'Japanese', None],
            ['Tongyi', 'qwen-turbo', 'Korean', None],
        ]
    )

    translator_gui.launch(share=True)